#include "EnclaveVault.h"
#include <sgx_cpuid.h>
#include "Vault_t.h"
#include "../EnclaveWrap/EnclaveWrap.h"
#include "../CryptoModule/Rsa.h"
#include "../CryptoModule/Des.h"

EnclaveVault::EnclaveVault(const std::string &id)
	: id_(id)
{
    //int cpuinfo[4] = { 0 };
    //sgx_cpuid(cpuinfo, 0);
	//
    //const int id_length = hex_string_.encrypt_size((const char *)cpuinfo, sizeof(cpuinfo));
    //id_ = new char[id_length + 1];
    //memset((void *)id_, 0, id_length + 1);
    //hex_string_.encrypt((const char *)cpuinfo, sizeof(cpuinfo), (char *)id_, id_length);

    log_info("EnclaveVault id: %s\n", id_.c_str());
}


EnclaveVault::~EnclaveVault()
{
    if (crypto_) {
        delete crypto_;
    }

    if (asymmetric_)
    {
        delete asymmetric_;
    }
}

int EnclaveVault::initEnclaveVault()
{
    int retval = 0;
    ocall_is_dir_exsit(&retval, user_data_dir_.c_str());
    if (retval == -1)
    {
        ocall_make_dir(&retval, user_data_dir_.c_str());
    }

    std::string valuts_dir = user_data_dir_ + "/" + id_;
    ocall_is_dir_exsit(&retval, valuts_dir.c_str());
    if (retval == -1)
    {
        ocall_make_dir(&retval, valuts_dir.c_str());
    }

    if (retval == -1)
    {
        log_error("make dir error");
        return -1;
    }

    initRsaSecurity();

    return initDesSecurity();
}

void EnclaveVault::initRsaSecurity()
{
    std::string pub_key;
    std::string pri_key;

    if (readFile(pub_key, pub_key_full_name()) == 0)
    {
		readFile(pri_key, pri_key_full_name());
        asymmetric_ = new Rsa(pub_key.c_str(), pri_key.c_str());
    }
    else
    {
        asymmetric_ = new Rsa(asymmetric_size_);

        writeFile(std::string(asymmetric_->public_key(), asymmetric_->public_key_length()), pub_key_full_name());
        writeFile(std::string(asymmetric_->private_key(), asymmetric_->private_key_length()), pri_key_full_name());
    }
}

int EnclaveVault::initDesSecurity()
{
	if (asymmetric_->public_key_length() != 0 && asymmetric_->private_key_length() == 0)
	{
		return 1;
	}

    std::string des_key_base64;

    if (readFile(des_key_base64, crypto_key_full_name()) == 0)
    {
        char *base64_decrypt_buf = new char[des_key_base64.length()];
        memset(base64_decrypt_buf, 0, des_key_base64.length());
        
        int len = base64_.decrypt(des_key_base64.c_str(), (int)des_key_base64.length(), base64_decrypt_buf, (int)des_key_base64.length());

        char *asymmetric_decrypt_buf = new char[len];
        memset(asymmetric_decrypt_buf, 0, len);
        len = asymmetric_->private_key_decrypt(base64_decrypt_buf, len, asymmetric_decrypt_buf, len);

        if (len == 0)
        {
            delete[]base64_decrypt_buf;
            delete[]asymmetric_decrypt_buf;
            return 1;
        }

        crypto_ = new Des(asymmetric_decrypt_buf, crypto_key_length_);

        delete[]base64_decrypt_buf;
        delete[]asymmetric_decrypt_buf;
    }
    else
    {
        crypto_ = new Des(crypto_key_length_);

        char *asymmetric_encrypt_buf = new char[asymmetric_size_];
        memset(asymmetric_encrypt_buf, 0, asymmetric_size_);
        asymmetric_->public_key_encrypt(crypto_->key(), crypto_key_length_, asymmetric_encrypt_buf, asymmetric_size_);

        char *base64_buf = new char[asymmetric_size_ * 2];
        memset(base64_buf, 0, asymmetric_size_ * 2);
        base64_.encrypt(asymmetric_encrypt_buf, asymmetric_size_, base64_buf, asymmetric_size_ * 2);

        des_key_base64 = base64_buf;
        writeFile(des_key_base64, crypto_key_full_name());

        delete[]asymmetric_encrypt_buf;
        delete[]base64_buf;
    }

	return 0;
}

int EnclaveVault::readFile(std::string &out, const std::string & file_name)
{
    int return_code = 0;

    _int64 file = 0;

    ocall_fopen(&return_code, &file, file_name.c_str(), "rb");
    if (file == NULL) {
        return 1;
    }

    ocall_fseek(&return_code, file, 0, 2);
    long len = 0;
    ocall_ftell(&len, file);

    ocall_fseek(&return_code, file, 0, 0);

    char *buf = new char[len + 1];
    memset(buf, 0, len + 1);

    size_t size_t_return_code = 0;
    ocall_fread(&size_t_return_code, file, buf, len + 1, len + 1, 1);

    ocall_fclose(&return_code, file);

    out = buf;

    delete[]buf;

    return 0;
}

int EnclaveVault::writeFile(std::string &in, const std::string & file_name)
{
    int return_code = 0;

    _int64 file = 0;
    ocall_fopen(&return_code, &file, file_name.c_str(), "wb");
    if (file == NULL)
    {
        assert(false);
        return 1;
    }

    sgx_status_t status = ocall_fputs(&return_code, file, in.c_str());
    if (status != SGX_SUCCESS)
    {
        return 1;
    }

    status = ocall_fclose(&return_code, file);
    if (status != SGX_SUCCESS)
    {
        return 1;
    }

    return 0;
}

int EnclaveVault::encrypt(const char * in_buf, int in_buf_len, char * out_buf, int out_buf_size)
{
    if (crypto_ == nullptr)
    {
        return 0;
    }

    int len = crypto_->encrypt(in_buf, in_buf_len, out_buf, out_buf_size);
    if (len == 0)
    {
        return 0;
    }

    int base64_buf_len = base64_.encrypt_size(len);
    char *base64_buf = new char[base64_buf_len];
	memset(base64_buf, 0, base64_buf_len);

    len = base64_.encrypt(out_buf, len, base64_buf, base64_buf_len);

    if (len == 0)
    {
        delete[]base64_buf;
        return 0;
    }

    memset(out_buf, 0, out_buf_size);
    memcpy(out_buf, base64_buf, len);

    delete[]base64_buf;

    return len;
}

int EnclaveVault::encrypt_size(int buf_len)
{
	if (crypto_ == nullptr)
	{
		return 0;
	}

	int len = crypto_->encrypt_size(buf_len);

	return base64_.encrypt_size(len);
}

int EnclaveVault::decrypt(const char * in_buf, int in_buf_len, char * out_buf, int out_buf_size)
{
    if (crypto_ == nullptr)
    {
        return 0;
    }

    int len = base64_.decrypt(in_buf, in_buf_len, out_buf, out_buf_size);
    if (len == 0)
    {
        return 0;
    }

    len = crypto_->decrypt(out_buf, len, out_buf, out_buf_size);

    memset(out_buf + len, 0, out_buf_size - len);

    return len;
}

int EnclaveVault::public_key_encrypt(const char * in_buf, int in_buf_len, char * out_buf, int out_buf_size)
{
    if (asymmetric_ == nullptr)
    {
        return 0;
    }

    int len = asymmetric_->public_key_encrypt(in_buf, in_buf_len, out_buf, out_buf_size);
    if (len == 0)
    {
        return 0;
    }

    int base64_buf_len = base64_.encrypt_size(len);
    char *base64_buf = new char[base64_buf_len];

    len = base64_.encrypt(out_buf, len, base64_buf, base64_buf_len);

    memset(out_buf, 0, out_buf_size);
    memcpy(out_buf, base64_buf, len);

    delete[]base64_buf;

    return len;
}

int EnclaveVault::private_key_decrypt(const char * in_buf, int in_buf_len, char * out_buf, int out_buf_size)
{
    if (asymmetric_ == nullptr)
    {
        return 0;
    }

    int len = base64_.decrypt(in_buf, in_buf_len, out_buf, out_buf_size);
    if (len == 0)
    {
        return 0;
    }

    len = asymmetric_->private_key_decrypt(out_buf, len, out_buf, out_buf_size);

    memset(out_buf + len, 0, out_buf_size - len);

    return len;
}

int EnclaveVault::private_key_encrypt(const char * in_buf, int in_buf_len, char * out_buf, int out_buf_size)
{
    if (asymmetric_ == nullptr)
    {
        return 0;
    }

    int len = asymmetric_->private_key_encrypt(in_buf, in_buf_len, out_buf, out_buf_size);
    if (len == 0)
    {
        return 0;
    }

    int base64_buf_len = base64_.encrypt_size(len);
    char *base64_buf = new char[base64_buf_len];

    len = base64_.encrypt(out_buf, len, base64_buf, base64_buf_len);

    memset(out_buf, 0, out_buf_size);
    memcpy(out_buf, base64_buf, len);

    delete[]base64_buf;

    return len;
}

int EnclaveVault::public_key_decrypt(const char * in_buf, int in_buf_len, char * out_buf, int out_buf_size)
{
    if (asymmetric_ == nullptr)
    {
        return 0;
    }

    int len = base64_.decrypt(in_buf, in_buf_len, out_buf, out_buf_size);
    if (len == 0)
    {
        return 0;
    }

    len = asymmetric_->public_key_decrypt(out_buf, len, out_buf, out_buf_size);

    memset(out_buf + len, 0, out_buf_size - len);

    return len;
}

int EnclaveVault::sign(const char * in_buf, int in_buf_len, char * out_buf, int out_buf_size)
{
	if (asymmetric_ == nullptr)
	{
		return 0;
	}

	int len = asymmetric_->private_key_encrypt(in_buf, in_buf_len, out_buf, out_buf_size);
	if (len == 0)
	{
		return 0;
	}

	int  hex_string_buf_len = hex_string_.encrypt_size(len);
	char *hex_string_buf = new char[hex_string_buf_len];

	len = hex_string_.encrypt(out_buf, len, hex_string_buf, hex_string_buf_len);

    if (len == 0)
    {
        delete[]hex_string_buf;
        return 0;
    }

	memset(out_buf, 0, out_buf_size);
	memcpy(out_buf, hex_string_buf, len);

	delete[]hex_string_buf;

	return len;
}

int EnclaveVault::sha256(const char * in_buf, int in_buf_len, char * out_buf, int out_buf_size)
{
	int len = sha256_.encrypt(in_buf, in_buf_len, out_buf, out_buf_size);
	if (len == 0)
	{
		return 0;
	}

	int  hex_string_buf_len = hex_string_.encrypt_size(len);
	char *hex_string_buf = new char[hex_string_buf_len];
	memset(hex_string_buf, 0, hex_string_buf_len);

	len = hex_string_.encrypt(out_buf, len, hex_string_buf, hex_string_buf_len);

    if (len == 0)
    {
        delete[]hex_string_buf;
        return 0;
    }

	memset(out_buf, 0, out_buf_size);
	memcpy(out_buf, hex_string_buf, len);

	delete[]hex_string_buf;

	return len;
}
